#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Author: bao
# @Date:   2017-01-23 18:50:01
# @Last Modified by:   bao
# @Last Modified time: 2017-05-15 23:11:38

from datetime import datetime
import sys
import copy
import math
import random
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
import demand_estimation as dm_est

from gen_topo_fully import *
from gen_topo_detour import *
from gen_topo_ovs import *

from flow_gen import *
from topo_gen import *


def sort_nonzeros1D(x):
    sidx = np.argsort(x)
    out_idx = sidx[np.in1d(sidx, np.flatnonzero(x!=0))][::-1]
    out_x = x[out_idx]
    return out_x, out_idx

def sort_nonzeros2D(x):
    x1 = np.where(x==0, np.nan, x)
    sidx = np.argsort(x1,1)[:,::-1]

    n = x.shape[1]
    extent_idx = (x==0).sum(1)
    valid_mask = extent_idx[:,None] <= np.arange(n)
    split_idx = (n-extent_idx[:-1]).cumsum()

    out_idx = np.split(sidx[valid_mask], split_idx)
    y = x[np.arange(x.shape[0])[:,None], sidx]
    out_x = np.split(y[valid_mask], split_idx)
    return out_x, out_idx

if __name__ == '__main__':
    
    random.seed(3)
    # sys paramters
    switch_nums = 8
    host_num = 6
    host_nums = switch_nums*host_num
    degree = 6
    row_size = int(math.ceil(math.sqrt(switch_nums)))
    threshold = 0.2
    demand_type = 'random' #cluster
    cluster_size = 5

    for i in range(1, len(sys.argv)):
        if sys.argv[i] == 'switch_nums':
            switch_nums = int(sys.argv[i+1])
        if sys.argv[i] == 'host_num':
            host_num = int(sys.argv[i+1])    
        if sys.argv[i] == 'degree':
            degree = int(sys.argv[i+1])
        if sys.argv[i] == 'demand_type':
            demand_type = sys.argv[i+1]
        if sys.argv[i] == 'row_size':
            row_size = sys.argv[i+1]
    
    # initial
    switchList = {}
    host_location = {} # {hostId: switchId}

    for switch in xrange(1,switch_nums+1):
        hosts = []
        for host in xrange(0,host_num):
            hosts.append(host+(switch-1)*host_num)
            host_location[host+(switch-1)*host_num]=(switch, degree+host+1)
        switchList[switch] = hosts
    
    # initial bipartite graph
    tf_graph = nx.Graph()
    wl_graph = nx.Graph()

    bp_src = range (1,switch_nums+1)
    bp_dst = range (switch_nums+1, switch_nums*2+1)

    for src in bp_src:
        tf_graph.add_node(src, bipartite=0, degree=0)
        wl_graph.add_node(src, bipartite=0, degree=0)

    for dst in bp_dst:
        tf_graph.add_node(dst, bipartite=1, degree=0)
        wl_graph.add_node(dst, bipartite=1, degree=0)

    # calculate initial matchings
    #print "k_regular_network"
    pre_matches = k_regular_network(switch_nums, degree)

    # 0.Make flows
    flows_src = {}
    flows_dst = {}
    makeFlows(flows_src, flows_dst, demand_type, switch_nums, host_num, host_nums, cluster_size, threshold)
  
    #-------------------------------------------------------------------
    #initial
    srcToRs = {} # {srcToR:{dstToR1:[portId1], dstToR2:[portId2}}
    dstToRs = {} # {dstToR:{srcToR1:[portId1], srcToR2:[portId2}}
    srcPorts = {} # {srcToR:{portId1:[dstToR1], portId2:[dstToR2}}
    dstPorts = {} # {dstToR:{portId1:[dstToR1], portId2:[dstToR2}}

    for portId in pre_matches:
        for (src, dst) in pre_matches[portId]:
            
            dst -= switch_nums
            print ((src, dst))
            #-----------------
            if src not in srcToRs:
                srcToRs[src] = {}
                srcPorts[src] = {}
                srcToRs[src][dst] = [portId]
                srcPorts[src][portId] = [dst]
            else:
                if dst not in srcToRs[src]:
                    srcToRs[src][dst] = [portId]
                else:
                    srcToRs[src][dst].append(portId)

                if portId not in srcToRs[src]:
                    srcPorts[src][portId] = [dst]
            #-----------------
            if dst not in dstToRs:
                dstToRs[dst] = {}
                dstPorts[dst] = {}
                dstToRs[dst][src] = [portId]
                dstPorts[dst][portId] = [src]

            else:
                if src not in dstToRs[dst]:
                    dstToRs[dst][src] = [portId]
                else:
                    dstToRs[dst][src].append(portId)

                if portId not in dstToRs[dst]:
                    dstToRs[dst][portId] = [src]
            
        print (srcToRs)
        print '----'
        print (dstToRs)
        print '----'
        print (srcPorts)
        print '----'
        print (dstPorts)
        print '-----------------'

    ucastToR = {}
    ucastPort = {}
    ucastToRFully = {}
    ucastPortFully = {}
    for torId in range(1, switch_nums+1):
        ucastToR[torId] = (range(1, degree+1), range(1, degree+1))
        ucastToRFully[torId] = (range(1, degree+1), range(1, degree+1))
    for portId in range(1, degree+1):
        ucastPort[portId] = (range(1, switch_nums+1), range(1, switch_nums+1))
        ucastPortFully[portId] = (range(1, switch_nums+1), range(1, switch_nums+1))
    
    starttime = datetime.now()
    # 1. flows demand estimation
    host_matrix2D = dm_est.demand_estimation(flows_src, flows_dst)
    # 2. ToR-level flows demand estimation
    sw_matrix2D = transform(host_matrix2D, switchList, host_num, demand_type, threshold=0.1)
    sw_matrix1D = sw_matrix2D.ravel()
    
    # 3. Greedy wavelegnth assignment
    weight_x, weight_idx = sort_nonzeros1D(sw_matrix1D)
    print weight_x
    print weight_idx
    print len(weight_x)
    
    for i in range(0, len(weight_idx)):
        srcToR = weight_idx[i]/switch_nums + 1
        dstToR = weight_idx[i]%switch_nums + 1 # + switch_nums
        print weight_idx[i], srcToR, dstToR
        if dstToR in srcToRs[srcToR]:
            print 'true', srcToR, dstToR
        else:
            temp_demand = host_num
            smallPortId = 1
            for portId in ucastToR[srcToR]:
                srcToR1 = dstToRs[portId]
                for x in xrange(1, 10):
                    pass


    endtime = datetime.now()
    timediff = endtime - starttime
    print ("Overall time: %s" %timediff)

    # Fully reconfigurable Non-blocking
    if False:
        # 3.generate topology and new matchings
        gen_topo_fully = gen_topo_fully()
        gen_topo_fully._transform(wl_graph, tf_graph, degree, switch_nums)
        cur_matches_fully = gen_topo_fully._matchings(wl_graph, pre_matches, degree, switch_nums)
        (d_matches, s_matches, new_matches) = gen_topo_fully._color_assignment(pre_matches, cur_matches_fully, switch_nums)
        # print len(d_matches), len(s_matches), len(new_matches)

        if False:
            for k in new_matches:
                pass
                # print 'new_matches',k, len(new_matches[k])
                # print sorted(new_matches[k])
            
            for k in s_matches:
                pass
                # print 's_matches',k, len(s_matches[k])
                # print sorted(s_matches[k])
            d_sum = 0
            for k in d_matches:
                d_sum += len(d_matches[k])
                # print 'd_matches',k, len(d_matches[k])
                # print sorted(d_matches[k])
    
    # Detour reconfigurable Non-blocking
    if False:
        # 3.generate topology and new matchings
        gen_topo_detour = gen_topo_detour()
        s_matches = {}
        d_matches = {}
        for x in pre_matches:
            s_matches[x] = copy.copy(pre_matches[x])
            d_matches[x] = []
        gen_topo_detour._wavelength_adjustment(pre_matches, d_matches, s_matches, row_size, switch_nums)
        gen_topo_detour._transform(wl_graph, tf_graph, degree, switch_nums)
        cur_matches_detour = gen_topo_detour._matchings(wl_graph, pre_matches, degree, switch_nums)
        (d_matches, s_matches, new_matches) = gen_topo_detour._color_assignment(pre_matches, cur_matches_detour, switch_nums)
        d_sum = 0
        for k in d_matches:
            d_sum += len(d_matches[k])
            # print 'd_matches',k, len(d_matches[k])
            # print sorted(d_matches[k])
        
        gen_topo_detour._wavelength_adjustment(new_matches, d_matches, s_matches, row_size, switch_nums)
        # print len(d_matches), len(s_matches), len(new_matches)

        if False:
            for k in new_matches:
                pass
                # print 'new_matches',k, len(new_matches[k])
                # print sorted(new_matches[k])
            
            for k in s_matches:
                pass
                # print 's_matches',k, len(s_matches[k])
                # print sorted(s_matches[k])
            d_sum = 0
            for k in d_matches:
                d_sum += len(d_matches[k])
                # print 'd_matches',k, len(d_matches[k])
                # print sorted(d_matches[k])
            # print d_sum
    
    # OvS reconfigurable with blocking
    if False:
        # 3.generate topology and new matchings
        gen_topo_ovs = gen_topo_ovs()
        gen_topo_ovs._transform(wl_graph, tf_graph, degree, row_size, switch_nums)
        cur_matches_fully = gen_topo_ovs._matchings(wl_graph, pre_matches, degree, switch_nums)
        (d_matches, s_matches, new_matches) = gen_topo_ovs._color_assignment(pre_matches, cur_matches_fully, switch_nums)
        # print len(d_matches), len(s_matches), len(new_matches)

        if False:

            for k in new_matches:
                pass
                # print 'new_matches',k, len(new_matches[k])
                # print sorted(new_matches[k])
            
            for k in s_matches:
                pass
                # print 's_matches',k, len(s_matches[k])
                # print sorted(s_matches[k])
            d_sum = 0
            for k in d_matches:
                d_sum += len(d_matches[k])
                # print 'd_matches',k, len(d_matches[k])
                # print sorted(d_matches[k])
            # print d_sum
